/*
  Copyright 2012 Sergey Ostanin

  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at

      http://www.apache.org/licenses/LICENSE-2.0

  Unless required by applicable law or agreed to in writing, software
  distributed under the License is distributed on an "AS IS" BASIS,
  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  See the License for the specific language governing permissions and
  limitations under the License.
*/

package ru.irenproject.store;

import java.io.ByteArrayInputStream;
import java.io.DataInput;
import java.io.DataInputStream;
import java.nio.file.Path;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map.Entry;
import java.util.Set;

import javax.annotation.Nullable;
import javax.inject.Inject;

import org.apache.commons.lang3.Validate;

import ru.irenproject.Oops;
import ru.irenproject.ProtoRegistry;

import com.almworks.sqlite4java.SQLiteConnection;
import com.almworks.sqlite4java.SQLiteException;
import com.almworks.sqlite4java.SQLiteStatement;
import com.google.common.io.ByteArrayDataOutput;
import com.google.common.io.ByteStreams;
import com.google.inject.assistedinject.Assisted;
import com.google.protobuf.ByteString;
import com.google.protobuf.Message;

public final class Store implements AutoCloseable {
  public static final ByteString ROOT_ID = Ids.ZERO;
  private static final int DEFAULT_MAX_UNPERSISTED_ACTIVITIES = 10;

  private final SecureRandom fRandom = new SecureRandom();

  private SQLiteConnection fConnection;
  private Checkpointer fCheckpointer;
  private boolean fInTransaction;
  private boolean fActive;
  private int fUnpersistedActivityCount;
  private final int fMaxUnpersistedActivities;

  private SQLiteStatement fGetThingSql;
  private int fGetThingIdParameter;
  private int fGetThingKindColumn;
  private int fGetThingValueColumn;
  private int fGetThingAspectsColumn;
  private int fGetThingListenersColumn;

  private SQLiteStatement fLookUpImmutableSql;
  private SQLiteStatement fInsertImmutableSql;

  private LinkedHashMap<ByteString, AbstractThing> fWorkspace;
  private LinkedHashSet<PostedEvent> fBus = new LinkedHashSet<>();
  private HashMap<ByteString, Message> fOriginals;

  private final ThingFactory fThingFactory;
  private final ProtoRegistry fProtoRegistry;

  private final MessageDigest fDigest;

  @Inject Store(ThingFactory thingFactory, ProtoRegistry protoRegistry,
      @Assisted @Nullable Path directory) {
    boolean ok = false;
    try {
      fThingFactory = thingFactory;
      fProtoRegistry = protoRegistry;
      fDigest = MessageDigest.getInstance("SHA-256");

      boolean inMemory = (directory == null);
      fMaxUnpersistedActivities = inMemory ? 1 : DEFAULT_MAX_UNPERSISTED_ACTIVITIES;

      Path dbFile = inMemory ? null : directory.resolve("db");
      fConnection = new SQLiteConnection(inMemory ? null : dbFile.toFile());
      fConnection.open();
      fConnection.exec("CREATE TABLE IF NOT EXISTS OBJECTS"
          + " (ID BLOB NOT NULL PRIMARY KEY, KIND INTEGER NOT NULL, VALUE BLOB NOT NULL"
          + ", LISTENERS BLOB, ASPECTS BLOB)");
      fConnection.exec("CREATE TABLE IF NOT EXISTS IMMUTABLES"
          + " (HASH BLOB NOT NULL, KIND INTEGER NOT NULL, ID BLOB NOT NULL"
          + ", PRIMARY KEY (HASH, KIND))");

      if (!inMemory) {
        fConnection.exec("PRAGMA journal_mode = WAL");
        fConnection.exec("PRAGMA wal_autocheckpoint = 50000");
        fConnection.exec("PRAGMA synchronous = NORMAL");

        fCheckpointer = Checkpointer.create(dbFile);
      }

      ok = true;
    } catch (SQLiteException | NoSuchAlgorithmException e) {
      throw new RuntimeException(e);
    } finally {
      if (!ok) {
        close();
      }
    }
  }

  @Override public void close() {
    if (fConnection != null) {
      fConnection.dispose();
      fConnection = null;
    }

    if (fCheckpointer != null) {
      fCheckpointer.finish();
      fCheckpointer = null;
    }
  }

  public ByteString generateId() {
    byte[] ba = new byte[Ids.SIZE];
    fRandom.nextBytes(ba);
    return Ids.fromBytes(ba);
  }

  public <T extends AbstractThing> T create(Class<T> cls) {
    Validate.isTrue(fActive);
    return createWithId(cls, generateId());
  }

  public <T extends AbstractThing> T createWithId(Class<T> cls, ByteString id) {
    Validate.isTrue(fActive);

    T res = fThingFactory.create(cls);
    res.setStore(this);
    res.setId(id);
    fWorkspace.put(id, res);

    res.initialize();
    return res;
  }

  public <T> T get(ByteString id, Class<T> cls) {
    return cls.cast(getAbstractThing(id));
  }

  public <T> List<T> getBatch(Iterable<ByteString> ids, Class<T> cls) {
    List<T> res = new ArrayList<>();
    getBatchTo(ids, cls, res);
    return res;
  }

  public <T> void getBatchTo(Iterable<ByteString> ids, Class<T> cls, Collection<? super T> out) {
    for (ByteString id : ids) {
      out.add(get(id, cls));
    }
  }

  public Thing getThing(ByteString id) {
    return getAbstractThing(id);
  }

  private AbstractThing getAbstractThing(ByteString id) {
    Validate.isTrue(fActive);

    try {
      AbstractThing res = fWorkspace.get(id);
      if (res == null) {
        fGetThingSql.reset();
        try {
          fGetThingSql.bind(fGetThingIdParameter, id.toByteArray());
          if (!fGetThingSql.step()) {
            throw Oops.format("Thing %s not found.", Ids.toHexString(id));
          }

          Kind kind = Kind.get(fGetThingSql.columnLong(fGetThingKindColumn));
          Class<? extends AbstractThing> thingClass = fProtoRegistry.getThingClass(kind);

          res = fThingFactory.create(thingClass);
          res.setStore(this);
          res.setId(id);

          byte[] value = fGetThingSql.columnBlob(fGetThingValueColumn);
          if (value == null) {
            value = new byte[0];
          }
          fThingFactory.loadStash(new ByteArrayInputStream(value), res);

          byte[] listeners = fGetThingSql.columnBlob(fGetThingListenersColumn);
          if (listeners != null) {
            loadListeners(listeners, res);
          }

          byte[] aspects = fGetThingSql.columnBlob(fGetThingAspectsColumn);
          if (aspects != null) {
            loadAspects(aspects, res);
          }
        } finally {
          fGetThingSql.reset();
        }

        if (fOriginals != null) {
          Validate.isTrue(fOriginals.put(id, fThingFactory.getStashMessage(res)) == null);
        }

        fWorkspace.put(id, res);
      }
      return res;
    } catch (SQLiteException e) {
      throw new RuntimeException(e);
    }
  }

  public <T extends AbstractThing> T obtainImmutable(Message message, Class<T> cls) {
    Validate.isTrue(fActive);

    Kind kind = fProtoRegistry.getThingKind(cls);
    ByteString hash = hash(message.toByteString());
    T res;

    ByteString id = lookUpImmutable(hash, kind);
    if (id == null) {
      res = create(cls);
      fThingFactory.setStashMessage(res, message);
      insertImmutable(hash, kind, res.id());
    } else {
      res = get(id, cls);
    }

    return res;
  }

  private ByteString lookUpImmutable(ByteString hash, Kind kind) {
    try {
      SQLiteStatement s = fLookUpImmutableSql;
      s.reset();
      try {
        s.bind(s.getBindParameterIndex(":HASH"), hash.toByteArray());
        s.bind(s.getBindParameterIndex(":KIND"), kind.value());
        return s.step() ? Ids.fromBytes(s.columnBlob(0)) : null;
      } finally {
        s.reset();
      }
    } catch (SQLiteException e) {
      throw new RuntimeException(e);
    }
  }

  private void insertImmutable(ByteString hash, Kind kind, ByteString id) {
    try {
      SQLiteStatement s = fInsertImmutableSql;
      s.reset();
      try {
        s.bind(s.getBindParameterIndex(":HASH"), hash.toByteArray());
        s.bind(s.getBindParameterIndex(":KIND"), kind.value());
        s.bind(s.getBindParameterIndex(":ID"), id.toByteArray());
        s.step();
      } finally {
        s.reset();
      }
    } catch (SQLiteException e) {
      throw new RuntimeException(e);
    }
  }

  private ByteString hash(ByteString data) {
    fDigest.reset();
    return ByteString.copyFrom(fDigest.digest(data.toByteArray()));
  }

  /**
   * @return ids of modified/created Thing-s
   */
  public Set<ByteString> dispatchEventsAndSave() {
    Validate.isTrue(fActive);
    Validate.notNull(fOriginals);

    dispatchEvents();
    return save();
  }

  private void dispatchEvents() {
    Iterator<PostedEvent> i;
    while ((i = fBus.iterator()).hasNext()) {
      PostedEvent e = i.next();
      i.remove();

      Thing source = getThing(e.sourceId());
      AbstractThing target = getAbstractThing(e.targetId());
      target.handleEvent(e.event(), source);
    }
  }

  public void getMarkedForRender(Set<ByteString> ids) {
    Validate.isTrue(fActive);

    for (AbstractThing t : fWorkspace.values()) {
      if (t.isMarkedForRender()) {
        ids.add(t.id());
      }
    }
  }

  public void beginActivity() {
    Validate.isTrue(!fActive);

    try {
      transactionNeeded();

      try {
        fWorkspace = new LinkedHashMap<>();
        fBus = new LinkedHashSet<>();
        fOriginals = new HashMap<>();

        fConnection.exec("SAVEPOINT ACTIVITY_START");
        fActive = true;
      } finally {
        if (!fActive) {
          cleanUpBeforeActivityEnd();
        }
      }
    } catch (SQLiteException e) {
      throw new RuntimeException(e);
    }
  }

  public void endActivity(boolean applyChanges) {
    Validate.isTrue(fActive);

    try {
      if (applyChanges) {
        Validate.isTrue(fOriginals == null);
        cleanUpBeforeActivityEnd();
        fConnection.exec("RELEASE ACTIVITY_START");

        fUnpersistedActivityCount = (fUnpersistedActivityCount + 1) % fMaxUnpersistedActivities;
        if (fUnpersistedActivityCount == 0) {
          persist();
        }

        if (fCheckpointer != null) {
          fCheckpointer.restartCountdown();
        }
      } else {
        cleanUpBeforeActivityEnd();
        fConnection.exec("ROLLBACK TO ACTIVITY_START");
        fConnection.exec("RELEASE ACTIVITY_START");
      }
    } catch (SQLiteException e) {
      throw new RuntimeException(e);
    }
  }

  private void transactionNeeded() {
    try {
      if (!fInTransaction) {
        fConnection.exec("BEGIN TRANSACTION");

        fGetThingSql = fConnection.prepare("SELECT KIND, VALUE, LISTENERS, ASPECTS FROM OBJECTS WHERE ID = :ID");
        fGetThingIdParameter = fGetThingSql.getBindParameterIndex(":ID");
        fGetThingKindColumn = getColumnIndex(fGetThingSql, "KIND");
        fGetThingValueColumn = getColumnIndex(fGetThingSql, "VALUE");
        fGetThingListenersColumn = getColumnIndex(fGetThingSql, "LISTENERS");
        fGetThingAspectsColumn = getColumnIndex(fGetThingSql, "ASPECTS");

        fLookUpImmutableSql = fConnection.prepare("SELECT ID FROM IMMUTABLES WHERE HASH = :HASH AND KIND = :KIND");
        fInsertImmutableSql = fConnection.prepare("INSERT INTO IMMUTABLES (HASH, KIND, ID)"
            + " VALUES (:HASH, :KIND, :ID)");

        fInTransaction = true;
      }
    } catch (SQLiteException e) {
      throw new RuntimeException(e);
    } finally {
      if (!fInTransaction) {
        cleanUpBeforeTransactionEnd();
        try {
          fConnection.exec("ROLLBACK");
        } catch (SQLiteException e) {}
      }
    }
  }

  public void persist() {
    Validate.isTrue(!fActive);

    try {
      if (fInTransaction) {
        cleanUpBeforeTransactionEnd();
        fConnection.exec("COMMIT");

        if (fCheckpointer != null) {
          fCheckpointer.advanceSequenceNumber();
        }
      }
    } catch (SQLiteException e) {
      throw new PersistException(e);
    }
  }

  private static int getColumnIndex(SQLiteStatement sql, String columnName) {
    try {
      int columnCount = sql.columnCount();
      for (int i = 0; i < columnCount; ++i) {
        if (columnName.equals(sql.getColumnName(i))) {
          return i;
        }
      }
      throw Oops.format("Column %s not found.", columnName);
    } catch (SQLiteException e) {
      throw new RuntimeException(e);
    }
  }

  private void cleanUpBeforeTransactionEnd() {
    if (fGetThingSql != null) {
      fGetThingSql.dispose();
      fGetThingSql = null;
    }

    if (fLookUpImmutableSql != null) {
      fLookUpImmutableSql.dispose();
      fLookUpImmutableSql = null;
    }

    if (fInsertImmutableSql != null) {
      fInsertImmutableSql.dispose();
      fInsertImmutableSql = null;
    }

    fInTransaction = false;
  }

  private void cleanUpBeforeActivityEnd() {
    fWorkspace = null;
    fBus = null;
    fOriginals = null;

    fActive = false;
  }

  /**
   * @return ids of modified/created Thing-s
   */
  private Set<ByteString> save() {
    try {
      Set<ByteString> res = new HashSet<>();

      SQLiteStatement insert = fConnection.prepare("INSERT INTO OBJECTS"
          + " (ID, KIND, VALUE, LISTENERS, ASPECTS)"
          + " VALUES (:ID, :KIND, :VALUE, :LISTENERS, :ASPECTS)");
      try {
        int insertId = insert.getBindParameterIndex(":ID");
        int insertKind = insert.getBindParameterIndex(":KIND");
        int insertValue = insert.getBindParameterIndex(":VALUE");
        int insertListeners = insert.getBindParameterIndex(":LISTENERS");
        int insertAspects = insert.getBindParameterIndex(":ASPECTS");

        SQLiteStatement update = fConnection.prepare("UPDATE OBJECTS SET VALUE = :VALUE WHERE ID = :ID");
        try {
          int updateValue = update.getBindParameterIndex(":VALUE");
          int updateId = update.getBindParameterIndex(":ID");

          SQLiteStatement update2 = fConnection.prepare("UPDATE OBJECTS SET"
              + " LISTENERS = :LISTENERS, ASPECTS = :ASPECTS WHERE ID = :ID");
          try {
            int update2Listeners = update2.getBindParameterIndex(":LISTENERS");
            int update2Aspects = update2.getBindParameterIndex(":ASPECTS");
            int update2Id = update2.getBindParameterIndex(":ID");

            for (AbstractThing t : fWorkspace.values()) {
              ByteString id = t.id();
              Message m = fThingFactory.getStashMessage(t);
              Message original = fOriginals.get(id);

              if (original == null) {
                insert.reset();
                insert.bind(insertId, id.toByteArray());
                insert.bind(insertKind, fProtoRegistry.getThingKind(t.getClass()).value());
                insert.bind(insertValue, m.toByteArray());
                insert.bind(insertListeners, nullIfEmpty(saveListeners(t)));
                insert.bind(insertAspects, nullIfEmpty(saveAspects(t)));
                insert.step();
                res.add(id);
              } else {
                if (!m.equals(original)) {
                  update.reset();
                  update.bind(updateValue, m.toByteArray());
                  update.bind(updateId, id.toByteArray());
                  update.step();
                  res.add(id);
                }
                if (t.listenersChanged() || t.aspectsChanged()) {
                  update2.reset();
                  update2.bind(update2Listeners, nullIfEmpty(saveListeners(t)));
                  update2.bind(update2Aspects, nullIfEmpty(saveAspects(t)));
                  update2.bind(update2Id, id.toByteArray());
                  update2.step();
                }
              }
            }
          } finally {
            update2.dispose();
          }
        } finally {
          update.dispose();
        }
      } finally {
        insert.dispose();
      }

      fOriginals = null;
      return res;
    } catch (SQLiteException e) {
      throw new RuntimeException(e);
    }
  }

  private static byte[] nullIfEmpty(byte[] value) {
    return (value.length > 0) ? value : null;
  }

  private static byte[] saveListeners(AbstractThing t) {
    ByteArrayDataOutput out = ByteStreams.newDataOutput();
    for (ByteString id : t.listenerIdList()) {
      Ids.write(id, out);
    }
    return out.toByteArray();
  }

  private static void loadListeners(byte[] data, AbstractThing t) {
    ByteArrayInputStream in = new ByteArrayInputStream(data);
    DataInput din = new DataInputStream(in);
    LinkedHashSet<ByteString> listeners = new LinkedHashSet<>();
    while (in.available() > 0) {
      listeners.add(Ids.read(din));
    }
    t.addListenersInternal(listeners);
  }

  private static byte[] saveAspects(AbstractThing t) {
    ByteArrayDataOutput out = ByteStreams.newDataOutput();
    for (Entry<AspectId, ByteString> e : t.aspectMap().entrySet()) {
      e.getKey().write(out);
      Ids.write(e.getValue(), out);
    }
    return out.toByteArray();
  }

  private static void loadAspects(byte[] data, AbstractThing t) {
    ByteArrayInputStream in = new ByteArrayInputStream(data);
    DataInput din = new DataInputStream(in);
    LinkedHashMap<AspectId, ByteString> aspects = new LinkedHashMap<>();
    while (in.available() > 0) {
      aspects.put(AspectId.read(din), Ids.read(din));
    }
    t.addAspectsInternal(aspects);
  }

  void post(ByteString sourceId, ByteString targetId, Event event) {
    fBus.add(new PostedEvent(sourceId, targetId, event));
  }
}
